#include "../include/SSS_smooth.h"
#include <cuda_runtime_api.h>
#include <driver_types.h>

#define gridsize 65535
#define blocksize 128

extern double pre_smooth_time ; 
extern double post_smooth_time ; 

void __global__  jacobi_kernel(const int i_1, const int i_n, const int s,int * d_row_ptr,int * d_col_idx,double *d_A_val, double * b, double * x)
{
    int N = abs(i_n - i_1) + 1;
    double w = 0.8;              //0.8
    int tid =blockIdx.x*blockDim.x+ threadIdx.x;

    double t;
    double d;

    int  j, k, begin_row, end_row;   
    
    if(tid<=i_n)
    {            
        t = b[tid];
        begin_row =d_row_ptr[tid], end_row = d_row_ptr[tid + 1];

        for (k = begin_row; k < end_row; ++k) 
        {
            j = d_col_idx[k];
            if (tid != j)
                t  -= d_A_val[k] * x[j];
            else if (tid == j)
            {
                d = d_A_val[k];
            }
        }
        if( d > SMALLFLOAT)
        {
            x[tid] = (1 - w) * x[tid] + w * t / d ;
        }
    }

}

static void jacobi_device(SSS_VEC * u, const int i_1, const int i_n, const int s,
    SSS_MAT * A, SSS_VEC * b, int L)
{
    while (L--) 
    {
        {
            cudaDeviceSynchronize();

            jacobi_kernel<<<gridsize,blocksize>>>(i_1, i_n,s, A->d_row_ptr,A->d_col_idx,A->d_val,b->d_d,u->d_d);
            cudaDeviceSynchronize();
        }
    }
}



void SSS_amg_smoother_pre_cuda(SSS_SMTR *s)
{
    //printf("pre_smooth \n");
    double pre_smooth_time_1 = SSS_get_time();

    SSS_MAT * A;
    SSS_VEC * b;
    SSS_VEC * x;

    int nsweeps;
    int istart;
    int iend;
    int istep;

    assert(s != NULL);

    A = s->A;
    b = s->b;
    x = s->x;

    nsweeps = s->nsweeps;
    istart = s->istart;
    iend = s->iend;
    istep = s->istep;

    jacobi_device(x, istart, iend, istep, A, b, nsweeps);   
    
    double pre_smooth_time_2 = SSS_get_time();
    pre_smooth_time+=pre_smooth_time_2-pre_smooth_time_1;

}

void SSS_amg_smoother_post_cuda(SSS_SMTR *s)
{
    //printf("post_smooth \n");

    double post_smooth_time_1 = SSS_get_time();

    SSS_MAT * A;
    SSS_VEC * b;
    SSS_VEC * x;

    int nsweeps;
    int istart;
    int iend;
    int istep;

    assert(s != NULL);

    A = s->A;
    b = s->b;
    x = s->x;
    nsweeps = s->nsweeps;
    istart = s->istart;
    iend = s->iend;
    istep = s->istep;

    jacobi_device(x, iend, istart, istep, A, b, nsweeps);

    double post_smooth_time_2 = SSS_get_time();
    post_smooth_time+=post_smooth_time_2-post_smooth_time_1;
}
